import torch
from torch.nn import Module
from torch.utils.data import DataLoader


class Client:
    validate_loader = None

    def __init__(self,
                 index: int,
                 n_clients: int,
                 dataloader: DataLoader,
                 model: Module,
                 lr: float = 0.1,
                 optimizer: str = "SGD",
                 weight_decay: float = 1e-4,
                 momentum: float = 0.9,
                 scheduler: str = "steplr",
                 step_size: int = 30,
                 gamma: float = 0.1,
                 compress_policy: str = None,
                 ) -> None:
        self.index = index
        self.n_clients = n_clients
        self.dataloader = dataloader
        self.model = model
        self.optimizer = None
        self.policy = compress_policy
        if optimizer == "SGD":
            self.optimizer = torch.optim.SGD(self.model.parameters(), lr=lr, weight_decay=weight_decay,
                                             momentum=momentum)
        else:
            pass

        if scheduler == "steplr":
            self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size, gamma)
        else:
            pass

    def save_weight(self):
        torch.save(self.model.state_dict(), f"./weight/client_side_weight_compress_{self.policy}.ckpt")

    def load_weight(self):
        self.model.load_state_dict(torch.load(f"./weight/client_side_weight_compress_{self.policy}.ckpt"))
